package scheduler

import (
	"context"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/sirupsen/logrus"

	"xagent/internal/metrics"
	"xagent/internal/types"
)

// TaskPriority represents the priority level of a task
type TaskPriority int

const (
	TaskPriorityLow      TaskPriority = 1
	TaskPriorityMedium   TaskPriority = 2
	TaskPriorityHigh     TaskPriority = 3
	TaskPriorityCritical TaskPriority = 4
)

// Priority levels for tasks
const (
	PriorityLow      = 1
	PriorityMedium   = 2
	PriorityHigh     = 3
	PriorityCritical = 4
)

// TaskState represents the state of a task
type TaskState string

const (
	TaskStatePending   TaskState = "pending"
	TaskStateScheduled TaskState = "scheduled"
	TaskStateRunning   TaskState = "running"
	TaskStateCompleted TaskState = "completed"
	TaskStateFailed    TaskState = "failed"
	TaskStateBlocked   TaskState = "blocked"
)

// Task represents a scheduled task
type Task struct {
	ID                   string
	Type                 string
	Priority             TaskPriority
	RequiredCapabilities []string
	Metadata             map[string]interface{}
	AssignedTo           string
	Status               TaskState
	StartTime            time.Time
	EstimatedDuration    time.Duration
	ActualDuration       time.Duration
}

// ResourcePool manages resource allocation
type ResourcePool struct {
	maxResources     int
	resources        map[string]bool
	mutex            sync.RWMutex
	maxLoadPerAgent  int
	maxMigrationCost float64
}

func NewResourcePool(maxResources int, maxLoadPerAgent int, maxMigrationCost float64) *ResourcePool {
	return &ResourcePool{
		maxResources:     maxResources,
		resources:        make(map[string]bool),
		maxLoadPerAgent:  maxLoadPerAgent,
		maxMigrationCost: maxMigrationCost,
	}
}

func (r *ResourcePool) GetResourceCount() int {
	r.mutex.RLock()
	defer r.mutex.RUnlock()
	return len(r.resources)
}

// Scheduler manages task scheduling and resource allocation
type Scheduler struct {
	tasks              map[string]*Task
	dependencies       map[string][]string
	resourcePool       *ResourcePool
	maxConcurrent      int
	running            map[string]bool
	taskQueue          []*Task
	logger             *logrus.Entry
	metrics            *SchedulerMetrics
	mutex              sync.RWMutex
	agents             map[string]types.Agent
	agentMetrics       map[string]*metrics.AgentMetrics
	workload           map[string]int
	collector          *metrics.Collector
	maxLoadPerAgent    int
	maxMigrationCost   float64
	batchStrategy      *BatchMigrationStrategy
	predictiveStrategy *PredictiveMigrationStrategy
	emergencyStrategy  *EmergencyMigrationStrategy
}

// SchedulerMetrics tracks scheduler performance
type SchedulerMetrics struct {
	TasksScheduled  int
	TasksCompleted  int
	TasksFailed     int
	AverageWaitTime time.Duration
	ResourceUsage   map[string]float64
	QueueLength     int
	mutex           sync.RWMutex
}

// TaskScore represents the suitability score of an agent for a task
type TaskScore struct {
	AgentID string
	Score   float64
}

// NewScheduler creates a new task scheduler
func NewScheduler(maxConcurrent int, resourcePool *ResourcePool) *Scheduler {
	return &Scheduler{
		tasks:              make(map[string]*Task),
		dependencies:       make(map[string][]string),
		resourcePool:       resourcePool,
		maxConcurrent:      maxConcurrent,
		running:            make(map[string]bool),
		taskQueue:          make([]*Task, 0),
		logger:             logrus.WithField("component", "scheduler"),
		metrics:            newSchedulerMetrics(),
		agents:             make(map[string]types.Agent),
		agentMetrics:       make(map[string]*metrics.AgentMetrics),
		workload:           make(map[string]int),
		maxLoadPerAgent:    resourcePool.maxLoadPerAgent,
		maxMigrationCost:   resourcePool.maxMigrationCost,
		batchStrategy:      NewBatchMigrationStrategy(nil, 5, time.Minute*5),
		predictiveStrategy: NewPredictiveMigrationStrategy(nil, 0.8, time.Minute*30),
		emergencyStrategy:  NewEmergencyMigrationStrategy(nil, 0.95),
	}
}

// newSchedulerMetrics creates new scheduler metrics
func newSchedulerMetrics() *SchedulerMetrics {
	return &SchedulerMetrics{
		ResourceUsage: make(map[string]float64),
	}
}

// AddTask adds a new task to the scheduler
func (s *Scheduler) AddTask(task *Task) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if task.ID == "" {
		task.ID = uuid.NewString()
	}

	if _, exists := s.tasks[task.ID]; exists {
		return fmt.Errorf("task with ID %s already exists", task.ID)
	}

	task.Status = TaskStatePending
	task.StartTime = time.Now()

	s.tasks[task.ID] = task
	s.updateDependencies(task)
	s.metrics.TasksScheduled++

	if s.canExecute(task) {
		s.addToQueue(task)
	}

	return nil
}

// updateDependencies updates the dependency graph
func (s *Scheduler) updateDependencies(task *Task) {
	if len(task.RequiredCapabilities) > 0 {
		s.dependencies[task.ID] = task.RequiredCapabilities
	}
}

// canExecute checks if a task can be executed
func (s *Scheduler) canExecute(task *Task) bool {
	if len(task.RequiredCapabilities) == 0 {
		return true
	}

	for _, depID := range task.RequiredCapabilities {
		if dep, exists := s.tasks[depID]; !exists || dep.Status != TaskStateCompleted {
			return false
		}
	}

	return true
}

// addToQueue adds a task to the execution queue
func (s *Scheduler) addToQueue(task *Task) {
	s.taskQueue = append(s.taskQueue, task)
	s.sortQueue()
}

// sortQueue sorts the task queue by priority
func (s *Scheduler) sortQueue() {
	sort.Slice(s.taskQueue, func(i, j int) bool {
		return s.taskQueue[i].Priority > s.taskQueue[j].Priority
	})
}

// Start starts the scheduler
func (s *Scheduler) Start(ctx context.Context) {
	go s.run(ctx)
}

// run is the main scheduler loop
func (s *Scheduler) run(ctx context.Context) {
	ticker := time.NewTicker(100 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			s.scheduleTasks()
		}
	}
}

// scheduleTasks attempts to schedule waiting tasks
func (s *Scheduler) scheduleTasks() {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if len(s.running) >= s.maxConcurrent {
		return
	}

	for i := 0; i < len(s.taskQueue); i++ {
		task := s.taskQueue[i]
		if s.canAllocateResources(task) {
			s.executeTask(task)
			s.taskQueue = append(s.taskQueue[:i], s.taskQueue[i+1:]...)
			i--
		}
	}
}

// canAllocateResources checks if required resources are available
func (s *Scheduler) canAllocateResources(task *Task) bool {
	s.resourcePool.mutex.RLock()
	defer s.resourcePool.mutex.RUnlock()

	return s.resourcePool.GetResourceCount() < s.resourcePool.maxResources
}

// allocateResources allocates resources for a task
func (s *Scheduler) allocateResources(task *Task) {
	s.resourcePool.mutex.Lock()
	defer s.resourcePool.mutex.Unlock()

	s.resourcePool.resources[task.ID] = true
}

// releaseResources releases resources allocated to a task
func (s *Scheduler) releaseResources(task *Task) {
	s.resourcePool.mutex.Lock()
	defer s.resourcePool.mutex.Unlock()

	delete(s.resourcePool.resources, task.ID)
}

// executeTask starts task execution
func (s *Scheduler) executeTask(task *Task) {
	s.running[task.ID] = true
	task.Status = TaskStateRunning
	task.StartTime = time.Now()

	s.allocateResources(task)

	go func() {
		defer s.taskCompleted(task)

		// Execute task
		if err := s.runTask(task); err != nil {
			s.handleTaskError(task, err)
			return
		}

		task.Status = TaskStateCompleted
		task.ActualDuration = time.Since(task.StartTime)
		s.metrics.TasksCompleted++
	}()
}

// runTask executes the actual task
func (s *Scheduler) runTask(task *Task) error {
	// TODO: Implement actual task execution
	time.Sleep(100 * time.Millisecond)
	return nil
}

// handleTaskError handles task execution errors
func (s *Scheduler) handleTaskError(task *Task, err error) {
	task.Status = TaskStateFailed
	s.metrics.TasksFailed++

	if task.Status == TaskStatePending {
		s.addToQueue(task)
	}

	s.logger.WithFields(logrus.Fields{
		"task_id": task.ID,
		"error":   err,
	}).Error("Task execution failed")
}

// taskCompleted handles task completion
func (s *Scheduler) taskCompleted(task *Task) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	delete(s.running, task.ID)
	s.releaseResources(task)

	// Check for dependent tasks
	for taskID, deps := range s.dependencies {
		for _, depID := range deps {
			if depID == task.ID {
				if dependentTask, exists := s.tasks[taskID]; exists {
					if s.canExecute(dependentTask) {
						s.addToQueue(dependentTask)
					}
				}
			}
		}
	}
}

// GetTaskStatus returns the status of a task
func (s *Scheduler) GetTaskStatus(taskID string) (*Task, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	task, exists := s.tasks[taskID]
	if !exists {
		return nil, fmt.Errorf("task with ID %s not found", taskID)
	}

	return task, nil
}

// GetMetrics returns scheduler metrics
func (s *Scheduler) GetMetrics() *SchedulerMetrics {
	s.metrics.mutex.RLock()
	defer s.metrics.mutex.RUnlock()

	s.metrics.QueueLength = len(s.taskQueue)
	s.metrics.ResourceUsage = make(map[string]float64)

	s.resourcePool.mutex.RLock()
	for resource, _ := range s.resourcePool.resources {
		used := float64(len(s.resourcePool.resources))
		s.metrics.ResourceUsage[resource] = used / float64(s.resourcePool.maxResources)
	}
	s.resourcePool.mutex.RUnlock()

	return s.metrics
}

// SetResourceCapacity sets the capacity for a resource
func (s *Scheduler) SetResourceCapacity(resource string, capacity float64) {
	s.resourcePool.mutex.Lock()
	defer s.resourcePool.mutex.Unlock()

	s.resourcePool.maxResources = int(capacity)
}

// GetResourceUsage returns current resource usage
func (s *Scheduler) GetResourceUsage() map[string]float64 {
	s.resourcePool.mutex.RLock()
	defer s.resourcePool.mutex.RUnlock()

	usage := make(map[string]float64)
	used := float64(len(s.resourcePool.resources))
	usage["total"] = used / float64(s.resourcePool.maxResources)
	return usage
}

// CancelTask cancels a running or pending task
func (s *Scheduler) CancelTask(taskID string) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	task, exists := s.tasks[taskID]
	if !exists {
		return fmt.Errorf("task with ID %s not found", taskID)
	}

	if task.Status == TaskStateRunning {
		// Release resources
		s.releaseResources(task)
		delete(s.running, task.ID)
	} else if task.Status == TaskStatePending {
		// Remove from queue
		for i, t := range s.taskQueue {
			if t.ID == taskID {
				s.taskQueue = append(s.taskQueue[:i], s.taskQueue[i+1:]...)
				break
			}
		}
	}

	delete(s.tasks, taskID)
	delete(s.dependencies, taskID)

	return nil
}

// UpdateTaskPriority updates the priority of a task
func (s *Scheduler) UpdateTaskPriority(taskID string, priority int) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	task, exists := s.tasks[taskID]
	if !exists {
		return fmt.Errorf("task with ID %s not found", taskID)
	}

	task.Priority = TaskPriority(priority)

	// Re-sort queue if task is pending
	if task.Status == TaskStatePending {
		s.sortQueue()
	}

	return nil
}

// GetQueuedTasks returns all queued tasks
func (s *Scheduler) GetQueuedTasks() []*Task {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	tasks := make([]*Task, len(s.taskQueue))
	copy(tasks, s.taskQueue)
	return tasks
}

// GetRunningTasks returns all currently running tasks
func (s *Scheduler) GetRunningTasks() []*Task {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	tasks := make([]*Task, 0, len(s.running))
	for taskID := range s.running {
		if task, exists := s.tasks[taskID]; exists {
			tasks = append(tasks, task)
		}
	}
	return tasks
}

// GetBlockedTasks returns all blocked tasks
func (s *Scheduler) GetBlockedTasks() []*Task {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	var blocked []*Task
	for _, task := range s.tasks {
		if task.Status == TaskStateBlocked {
			blocked = append(blocked, task)
		}
	}
	return blocked
}

// GetTaskDependencies returns the dependencies for a task
func (s *Scheduler) GetTaskDependencies(taskID string) ([]string, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	deps, exists := s.dependencies[taskID]
	if !exists {
		return nil, fmt.Errorf("task with ID %s not found", taskID)
	}

	return append([]string{}, deps...), nil
}

// AddAgent adds an agent to the scheduler
func (s *Scheduler) AddAgent(agent types.Agent) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if agent == nil {
		return fmt.Errorf("cannot add nil agent")
	}

	if _, exists := s.agents[agent.GetID()]; exists {
		return fmt.Errorf("agent %s already exists", agent.GetID())
	}

	s.agents[agent.GetID()] = agent
	s.workload[agent.GetID()] = 0
	return nil
}

// RemoveAgent removes an agent from the scheduler
func (s *Scheduler) RemoveAgent(agentID string) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if _, exists := s.agents[agentID]; !exists {
		return fmt.Errorf("agent %s not found", agentID)
	}

	delete(s.agents, agentID)
	delete(s.workload, agentID)
	delete(s.agentMetrics, agentID)
	return nil
}

// GetAgent returns an agent by ID
func (s *Scheduler) GetAgent(agentID string) (types.Agent, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	agent, exists := s.agents[agentID]
	if !exists {
		return nil, fmt.Errorf("agent %s not found", agentID)
	}

	return agent, nil
}

// SelectAgent selects the best agent for a task
func (s *Scheduler) SelectAgent(ctx context.Context, task *Task) (string, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	if len(s.agents) == 0 {
		return "", fmt.Errorf("no agents available")
	}

	var bestScore float64 = -1
	var bestAgent string

	for id, agent := range s.agents {
		score := s.calculateAgentScore(id, agent, task)
		if score > bestScore {
			bestScore = score
			bestAgent = id
		}
	}

	if bestAgent == "" {
		return "", fmt.Errorf("no suitable agent found")
	}

	return bestAgent, nil
}

// calculateAgentScore calculates a score for an agent based on various factors
func (s *Scheduler) calculateAgentScore(agentID string, agent types.Agent, task *Task) float64 {
	// Base score starts at 1.0
	score := 1.0

	// Factor 1: Current workload (0.0-1.0)
	workload := float64(s.workload[agentID]) / float64(s.maxLoadPerAgent)
	score *= (1.0 - workload)

	// Factor 2: Capability match (0.0-1.0)
	capScore := s.calculateCapabilityMatch(agent, task.RequiredCapabilities)
	score *= capScore

	// Factor 3: Agent state penalty
	if agent.GetState() != "idle" {
		score *= 0.5
	}

	return score
}

// calculateCapabilityMatch calculates how well an agent's capabilities match task requirements
func (s *Scheduler) calculateCapabilityMatch(agent types.Agent, taskCaps []string) float64 {
	if len(taskCaps) == 0 {
		return 1.0
	}

	agentCaps := agent.GetCapabilities()
	matches := 0

	for _, taskCap := range taskCaps {
		for _, agentCap := range agentCaps {
			if taskCap == agentCap {
				matches++
				break
			}
		}
	}

	return float64(matches) / float64(len(taskCaps))
}

// UpdateAgentMetrics updates metrics for an agent
func (s *Scheduler) UpdateAgentMetrics(agentID string, metrics *metrics.AgentMetrics) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.agentMetrics[agentID] = metrics
}

// UpdateWorkload updates the workload count for an agent
func (s *Scheduler) UpdateWorkload(agentID string, delta int) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.workload[agentID] += delta
	if s.workload[agentID] < 0 {
		s.workload[agentID] = 0
	}
}

// GetAgentLoad returns the current load for an agent
func (s *Scheduler) GetAgentLoad(agentID string) int {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	return s.workload[agentID]
}

// GetSystemLoad returns overall system load metrics
func (s *Scheduler) GetSystemLoad() map[string]interface{} {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	totalLoad := 0
	maxLoad := 0
	agentLoads := make(map[string]int)

	for agentID, load := range s.workload {
		totalLoad += load
		if load > maxLoad {
			maxLoad = load
		}
		agentLoads[agentID] = load
	}

	return map[string]interface{}{
		"total_load":    totalLoad,
		"max_load":      maxLoad,
		"average_load":  float64(totalLoad) / float64(len(s.workload)),
		"agent_loads":   agentLoads,
		"agent_count":   len(s.agents),
		"total_agents":  len(s.workload),
		"active_agents": len(s.agents),
	}
}

// MigrateTask migrates a task from one agent to another
func (s *Scheduler) MigrateTask(taskID, fromAgent, toAgent string) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	task, ok := s.tasks[taskID]
	if !ok {
		return ErrTaskNotFound
	}

	cost := s.calculateMigrationCost(task, fromAgent, toAgent)
	if cost > s.maxMigrationCost {
		return ErrMigrationTooExpensive
	}

	if !s.checkAgentCapacity(toAgent, task) {
		return ErrInsufficientCapacity
	}

	err := s.recordMigrationMetrics(taskID, fromAgent, toAgent, cost)
	if err != nil {
		return err
	}

	task.AssignedTo = toAgent
	s.workload[fromAgent]--
	s.workload[toAgent]++

	return nil
}

// checkAgentCapacity checks if an agent can handle a task
func (s *Scheduler) checkAgentCapacity(agentID string, task *Task) bool {
	agent := s.agents[agentID]
	if agent == nil {
		return false
	}

	// Check current load
	currentLoad := s.workload[agentID]
	if currentLoad >= s.maxLoadPerAgent {
		return false
	}

	// Check capabilities
	agentCaps := agent.GetCapabilities()
	for _, required := range task.RequiredCapabilities {
		found := false
		for _, available := range agentCaps {
			if required == available {
				found = true
				break
			}
		}
		if !found {
			return false
		}
	}

	return true
}

// recordMigrationMetrics records metrics about a task migration
func (s *Scheduler) recordMigrationMetrics(taskID, fromAgent, toAgent string, cost float64) error {
	metrics := map[string]interface{}{
		"task_id":    taskID,
		"from_agent": fromAgent,
		"to_agent":   toAgent,
		"total_cost": cost,
		"timestamp":  time.Now(),
	}

	err := s.collector.RecordEvent("task_migration", metrics)
	if err != nil {
		return err
	}

	return nil
}

// GetMigrationCandidates finds tasks that could be migrated to balance load
func (s *Scheduler) GetMigrationCandidates() []*TaskMigration {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	var candidates []*TaskMigration
	systemLoad := s.GetSystemLoad()
	avgLoad := systemLoad["average_load"].(float64)

	// Find overloaded agents
	for agentID, load := range s.workload {
		if float64(load) > avgLoad*1.2 { // 20% above average
			// Find tasks that could be migrated
			tasks := s.getAgentTasks(agentID)
			for _, task := range tasks {
				// Find potential target agents
				for targetID, targetLoad := range s.workload {
					if float64(targetLoad) < avgLoad*0.8 { // 20% below average
						cost := s.calculateMigrationCost(task, agentID, targetID)
						if cost < s.maxMigrationCost {
							candidates = append(candidates, &TaskMigration{
								TaskID:    task.ID,
								FromAgent: agentID,
								ToAgent:   targetID,
								Priority:  TaskPriority(task.Priority),
								Cost:      cost,
							})
						}
					}
				}
			}
		}
	}

	// Sort candidates by cost/benefit ratio
	sort.Slice(candidates, func(i, j int) bool {
		return candidates[i].Cost < candidates[j].Cost
	})

	return candidates
}

// getAgentTasks returns all tasks assigned to an agent
func (s *Scheduler) getAgentTasks(agentID string) []*Task {
	var tasks []*Task
	for _, task := range s.tasks {
		if task.AssignedTo == agentID {
			tasks = append(tasks, task)
		}
	}
	return tasks
}

// PlanBatchMigration plans a batch migration
func (s *Scheduler) PlanBatchMigration(ctx context.Context) (*MigrationPlan, error) {
	return s.batchStrategy.PlanBatchMigration(ctx)
}

// ExecuteBatchMigration executes a batch migration plan
func (s *Scheduler) ExecuteBatchMigration(ctx context.Context, plan *MigrationPlan) error {
	return s.batchStrategy.ExecuteBatchMigration(ctx, plan)
}

// PlanPredictiveMigration plans a predictive migration
func (s *Scheduler) PlanPredictiveMigration(ctx context.Context) (*MigrationPlan, error) {
	return s.predictiveStrategy.PlanPredictiveMigration(ctx)
}

// HandleEmergency handles an emergency situation
func (s *Scheduler) HandleEmergency(ctx context.Context, agentID string) error {
	return s.emergencyStrategy.HandleEmergency(ctx, agentID)
}

// calculateMigrationCost calculates the migration cost between two tasks
func (s *Scheduler) calculateMigrationCost(task *Task, fromAgent, toAgent string) float64 {
	// Implementation of migration cost calculation logic
	return 0.0
}

// GetCapabilities returns the task's required capabilities
func (t *Task) GetCapabilities() []string {
	if t.RequiredCapabilities != nil {
		return t.RequiredCapabilities
	}
	return []string{}
}

// GetAssignedTo returns the agent ID this task is assigned to
func (t *Task) GetAssignedTo() string {
	if t.AssignedTo != "" {
		return t.AssignedTo
	}
	return ""
}
